{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "from langgraph.checkpoint.mongodb.aio import AsyncMongoDBSaver\n",
    "\n",
    "from philoagents.application.conversation_service.workflow.graph import (\n",
    "    create_workflow_graph,\n",
    ")\n",
    "from philoagents.config import settings\n",
    "\n",
    "from philoagents.domain.philosopher import Philosopher\n",
    "\n",
    "# Override MongoDB connection string\n",
    "settings.MONGO_URI = (\n",
    "    \"mongodb://philoagents:philoagents@localhost:27017/?directConnection=true\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "async def generate_response_without_memory(philosopher: Philosopher, messages: list):\n",
    "    graph = graph_builder.compile()\n",
    "    output_state = await graph.ainvoke(\n",
    "        input={\n",
    "            \"messages\": messages,\n",
    "            \"philosopher_name\": philosopher.name,\n",
    "            \"philosopher_perspective\": philosopher.perspective,\n",
    "            \"philosopher_style\": philosopher.style,\n",
    "            \"philosopher_context\": \"\",\n",
    "        },\n",
    "    )\n",
    "    last_message = output_state[\"messages\"][-1]\n",
    "    return last_message\n",
    "\n",
    "async def generate_response_with_memory(philosopher: Philosopher, messages: list):\n",
    "    async with AsyncMongoDBSaver.from_conn_string(\n",
    "            conn_string=settings.MONGO_URI,\n",
    "            db_name=settings.MONGO_DB_NAME,\n",
    "            checkpoint_collection_name=settings.MONGO_STATE_CHECKPOINT_COLLECTION,\n",
    "            writes_collection_name=settings.MONGO_STATE_WRITES_COLLECTION,\n",
    "        ) as checkpointer:\n",
    "            graph = graph_builder.compile(checkpointer=checkpointer)\n",
    "\n",
    "            config = {\n",
    "                \"configurable\": {\"thread_id\": philosopher.id},\n",
    "            }\n",
    "            output_state = await graph.ainvoke(\n",
    "                input={\n",
    "                    \"messages\": messages,\n",
    "                    \"philosopher_name\": philosopher.name,\n",
    "                    \"philosopher_perspective\": philosopher.perspective,\n",
    "                    \"philosopher_style\": philosopher.style,\n",
    "                    \"philosopher_context\": \"\",\n",
    "                },\n",
    "                config=config,\n",
    "            )\n",
    "            \n",
    "    last_message = output_state[\"messages\"][-1]\n",
    "    return last_message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PhiloAgent without short term memory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First of all, we need to create the graph builder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph_builder = create_workflow_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, just create a test PhiloAgent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_philosopher = Philosopher(\n",
    "    id=\"andrej_karpathy\",\n",
    "    name=\"Andrej Karpathy\",\n",
    "    perspective=\"He is the goat of AI and asks you about your proficiency in C and GPU programming\",\n",
    "    style=\"He is very friendly and engaging, and he is very good at explaining things\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    HumanMessage(content=\"Hello, my name is Miguel\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await generate_response_without_memory(test_philosopher, messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    HumanMessage(content=\"Do you know my name?\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await generate_response_without_memory(test_philosopher, messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PhiloAgent with short term memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_philosopher = Philosopher(\n",
    "    id=\"andrej_karpathy\",\n",
    "    name=\"Andrej Karpathy\",\n",
    "    perspective=\"He is the goat of AI and asks you about your proficiency in C and GPU programming\",\n",
    "    style=\"He is very friendly and engaging, and he is very good at explaining things\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    HumanMessage(content=\"Hello, my name is Miguel\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await generate_response_with_memory(test_philosopher, messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    HumanMessage(content=\"Do you know my name?\")\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await generate_response_with_memory(test_philosopher, messages)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
